#!/usr/bin/env python3
#usage python anno_dir.py <video_directory_root>
#for mine_drill project, <video_directory_root> should named as '_'.join([video_month,camera_num,serial_num]), such as 201906_15113_3

import cv2
import numpy as np
from operator import mod
from math import ceil
import sys,os,time
from pathlib import Path
from tqdm import tqdm

try:
    mediaroot = sys.argv[1]
except:
    mediaroot = 'add-drill'

try:
    rootPath = Path(mediaroot).expanduser()
except:
    rootPath = Path(mediaroot)
filetype = '*.mp4'
mediafiles = rootPath.glob(filetype)
mediafiles = [str(mediafile) for mediafile in mediafiles]
mediafiles = sorted(mediafiles)
window_position_x = 1 
window_position_y = 1 


window_width = 640
window_height = 480

white_zero = cv2.imread('template/white_0.jpg')
white_one = cv2.imread('template/white_1.jpg')
white_two = cv2.imread('template/white_2.jpg')
white_three = cv2.imread('template/white_3.jpg')
white_four = cv2.imread('template/white_4.jpg')
white_five = cv2.imread('template/white_5.jpg')
white_six = cv2.imread('template/white_6.jpg')
white_seven = cv2.imread('template/white_7.jpg')
white_eight = cv2.imread('template/white_8.jpg')
white_night = cv2.imread('template/white_9.jpg')
black_zero = cv2.imread('template/black_0.jpg')
black_one = cv2.imread('template/black_1.jpg')
black_two = cv2.imread('template/black_2.jpg')
black_three = cv2.imread('template/black_3.jpg')
black_four = cv2.imread('template/black_4.jpg')
black_five = cv2.imread('template/black_5.jpg')
black_six = cv2.imread('template/black_6.jpg')
black_seven = cv2.imread('template/black_7.jpg')
black_eight = cv2.imread('template/black_8.jpg')
black_night = cv2.imread('template/black_9.jpg')
nums = [white_zero,white_one,white_two,white_three,white_four,white_five,white_six,white_seven,white_eight,white_night,black_zero,black_one,black_two,black_three,black_four,black_five,black_six,black_seven,black_eight,black_night]

def timing(position):
    current = position//25
    current_sec = mod(current,60)
    current_min = current // 60
    if current_min >= 60:
        current_min = mod(current_min, 60)
    current_hour = current // 3600
    current_time = ''.join([str(current_hour).zfill(2),':',str(current_min).zfill(2),':',str(current_sec).zfill(2)])

    return current_time

def get_last_anno():
    with open(annofile,'r') as f:
        tags = f.read().splitlines()
    index = -1
    if tags[index].isalpha():
        index -= 1 
    end_anno = tags[index]
    if '_' in end_anno:
        end_anno = 'NULL'

    return end_anno

def load_tag_state():
    with open(annofile,'r') as f:
        tags = f.read().splitlines()
    tag_states = [line for line in tags if line.isalpha()]
    try:
        tag_state_current = tag_states[-1]
    except:
        tag_state_current = 'None'

    return tag_state_current

def get_position(time):
    try:
        current_hour,current_min,current_sec = time.split(':')
        current_frame = int(fps * (int(current_hour) * 3600 + int(current_min) * 60 + int(current_sec)))
    except:
        current_frame = position

    return current_frame

def get_number(img):
    mean,stddv = cv2.meanStdDev(img)
    *_,channels = img.shape
    #imgMean = 121
    imgMean = sum(mean)/channels
    _,num_img = cv2.threshold(img,imgMean,255,cv2.THRESH_BINARY)
    #_, white_num =cv2.threshold(num_img,121,255,cv2.THRESH_BINARY)
    #inverse_img = 255 - num_img
    #_, black_num =cv2.threshold(inverse_img,121,255,cv2.THRESH_BINARY)
    #new_num_img = cv2.add(white_num,black_num)
    results = [cv2.matchTemplate(img,template,cv2.TM_CCOEFF_NORMED).astype(float) for template in nums]

    return mod(results.index(max(results)),10)

serial_num = 0
total_num = len(mediafiles)
print(''.join(["Total media: ",str(total_num)]))

for mediafile in mediafiles:
    serial_num += 1
    annofile = ''.join([mediafile,'.anno'])
    try:
        media = cv2.VideoCapture(mediafile)
    except:
        break
    fps = media.get(cv2.CAP_PROP_FPS)
    frames = int(media.get(cv2.CAP_PROP_FRAME_COUNT))
    if frames < 50:
        break
    total_time = timing(frames)
    loop_flag = 0
    clip_flag = 0
    edit_flag = 0
    position = 0

    try:
        _,image = media.read()
    except:
        break

    _,img = cv2.threshold(image,121,255,cv2.THRESH_BINARY)
    shape = image.shape
    height = shape[0]
    width = shape[1]

    time_font = cv2.FONT_HERSHEY_COMPLEX
    tag_font = cv2.FONT_HERSHEY_SIMPLEX
    help_font = cv2.FONT_ITALIC
    help_info_1 = 'help:  "a"/add; "r"/retrieve.; "u"/useless; "k"/ok; "="/next'
    help_info_2 = '      <Space> to tag start; <Enter> to tag end;"." to restore '
    help_info_3 = '      ESC/quit, p(ause), d/f/g(forward), v/b/n(back)'
    time_color = (209,80,0,255)
    tag_state_color = (209,80,0,255)
    help_color_1 = (109,80,50,155)
    help_color_2 = (109,80,50,155)

    if height == 720:
        font_size = 1
        time_position = (50,30)
        time_video_start_position = (40,140)
        font_size_time_video_start = 2
        font_bold_time_video_start = 2
        time_bold = 2
        tag_state_position = (400,30)
        tag_state_bold = 2
        tag_position = (600,30)
        tag_color = (209,80,0,255)
        tag_bold = 2
        help_position_1 = (50,630)
        help_bold_1 = 2
        help_position_2 = (50,710)
        help_bold_2 = 2
        help_position_3 = (50,670)
        help_bold_3 = 2
        serial_position = (50,600)
    elif height==1080:
        font_size = 2
        proportion = 1080/720
        time_position = (ceil(50*proportion),ceil(30*proportion))
        time_video_start_position = (ceil(50*proportion),ceil(630*proportion))
        time_bold = 2
        tag_state_position = (ceil(200+400*proportion),ceil(30*proportion))
        tag_state_bold = 2
        tag_position = (ceil(200+600*proportion),ceil(30*proportion))
        tag_color = (209,80,0,255)
        tag_bold = 2
        help_position_1 = (ceil(50*proportion),ceil(640*proportion))
        help_bold_1 = 2
        help_position_2 = (ceil(50*proportion),ceil(710*proportion))
        help_bold_2 = 2
        serial_position = (ceil(50*proportion),ceil(600*proportion))
        help_position_3 = (ceil(50*proportion),ceil(670*proportion))
        help_bold_3 = 2

    if height != 720:
        img = cv2.resize(img,(1280,720),interpolation=cv2.INTER_AREA)
    try:
        year_1 = str(get_number(img[53:85, 47:70]))
        year_2 = str(get_number(img[53:85, 71:94]))
        year_3 = str(get_number(img[53:85, 96:119]))
        year_4 = str(get_number(img[53:85, 119:142]))
        month_1 = str(get_number(img[53:85, 167:190]))
        month_2 = str(get_number(img[53:85, 191:214]))
        day_1 = str(get_number(img[53:85, 240:263]))
        day_2 = str(get_number(img[53:85, 263:286]))
        hour_1 = str(get_number(img[53:85, 479:502]))
        hour_2 = str(get_number(img[53:85, 503:526]))
        min_1 = str(get_number(img[53:85, 551:574]))
        min_2 = str(get_number(img[53:85, 575:598]))
        sec_1 = str(get_number(img[53:85, 623:646]))
        sec_2 = str(get_number(img[53:85, 647:670]))
        time_video_start = ''.join([year_1,year_2,year_3,year_4,month_1,month_2,day_1,day_2,'_',hour_1,hour_2,':',min_1,min_2,':',sec_1,sec_2])
    except:
        time_video_start = '00000000_00:00:00'
        break

    print(time_video_start)
    if not os.path.isfile(annofile):
        with open(annofile,'a') as f:
            f.writelines(''.join([time_video_start,'\n']));
    else:
        with open(annofile,'r') as f:
            anno_len = len(f.read().splitlines())
        if anno_len == 0:
            with open(annofile,'a') as f:
                f.writelines(''.join([time_video_start,'\n']));
            

    medianame = mediafile.split('/')[-1]
    print(medianame)
    print('fps = ', fps)
    total_s = media.get(cv2.CAP_PROP_FRAME_COUNT)
    print("total_s = ", total_s)

    title = ''.join([medianame,'\tESC(quit), p(ause), d/f/g(forward), v/b/n(back), x/z(oom),","(load last.'])
    cv2.namedWindow(title, 0)
    cv2.moveWindow(title,window_position_x,window_position_y)
    cv2.resizeWindow(title, window_width, window_height)
    cv2.createTrackbar('frame',title,0,frames,timing)


    frame_index = 0
    if media.isOpened():

        while True:
            if loop_flag == position:
                loop_flag = loop_flag + 1
                cv2.setTrackbarPos('frame',title,loop_flag)
            else:
                position = cv2.getTrackbarPos('frame',title)
                loop_flag = position
                media.set(cv2.CAP_PROP_POS_FRAMES,position)

            current_time = timing(position)
            time_info = ''.join([current_time,'/',total_time])

            tag_state = load_tag_state()
            try:
                last_line = get_last_anno()
            except:
                last_line = 'NULL'
                
            try:
                last_start,last_end = last_line.split('|')
            except:
                last_start = last_line.split('|')[0]
            if last_start != 'NULL':
                last_start_frame = get_position(last_start)
            else:
                last_start_frame = position
            if last_start_frame != position:
                clip_flag = 1
                tag_state = load_tag_state()

            wait_key = cv2.waitKeyEx(1)
            if wait_key == 112: # p
                cv2.waitKey(0)
            elif wait_key == 101: # e
                cmd = ''.join(["vim ",annofile])
                os.system(cmd)
                cv2.waitKey(0)
            elif wait_key == 103: # g
                position = loop_flag + 250
                loop_flag = position
                media.set(cv2.CAP_PROP_POS_FRAMES,position)
            elif wait_key == 102: # f
                position = loop_flag + 125
                loop_flag = position
                media.set(cv2.CAP_PROP_POS_FRAMES,position)
            elif wait_key == 100: # d
                position = loop_flag + 25
                loop_flag = position
                media.set(cv2.CAP_PROP_POS_FRAMES,position)
            elif wait_key == 98: # b
                position = loop_flag - 100
                loop_flag = position
                media.set(cv2.CAP_PROP_POS_FRAMES,position)
            elif wait_key == 118: # v
                position = loop_flag - 25
                loop_flag = position
                media.set(cv2.CAP_PROP_POS_FRAMES,position)
            elif wait_key == 110: # n
                position = loop_flag - 250
                loop_flag = position
                media.set(cv2.CAP_PROP_POS_FRAMES,position)
            elif wait_key == 108: # l
                position = frames - 100
                loop_flag = position
                media.set(cv2.CAP_PROP_POS_FRAMES,position)
            elif wait_key == 122: # z
                window_width = ceil(window_width * 1.2)
                window_height = ceil(window_height * 1.2)
                window_zoom = (window_width,window_height)
                cv2.resizeWindow(title, window_zoom)
            elif wait_key == 120: # x
                window_width = ceil(window_width * 0.8)
                window_height = ceil(window_height * 0.8)
                window_zoom = (window_width,window_height)
                cv2.resizeWindow(title, window_zoom)
            elif wait_key == 8: # backspace
                cmd = ''.join(["sed -i '$d' ",annofile])
                os.system(cmd)
            elif wait_key == 27: # ESC
                cv2.destroyAllWindows()
                sys.exit()
                break
            elif wait_key == 97: # a
                edit_flag = 1
                tag_state = 'add'
                with open(annofile,'a') as f:
                    f.writelines('add\n');
            elif wait_key == 114: # r
                edit_flag = 1
                tag_state = 'retrieve'
                with open(annofile,'a') as f:
                    f.writelines('retrieve\n');
            elif wait_key == 117: # u
                edit_flag = 1
                tag_state = 'xx'
                with open(annofile,'a') as f:
                    f.writelines('xx\n');
                position = frames
                loop_flag = position
                media.set(cv2.CAP_PROP_POS_FRAMES,position)
            elif wait_key == 107: # k
                edit_flag = 1
                tag_state = 'ok'
                with open(annofile,'a') as f:
                    f.writelines('ok\n');
                position = frames
                loop_flag = position
                media.set(cv2.CAP_PROP_POS_FRAMES,position)
            elif wait_key == 32: # Space
                clip_flag = 1
                edit_flag = 1
                current_frame = cv2.getTrackbarPos('frame',title)
                start_time = timing(current_frame - 10)
                with open(annofile,'a') as f:
                    f.writelines(''.join([start_time,'|']));
            elif wait_key == 13: # Enter
                clip_flag = 1
                edit_flag = 1
                current_frame = cv2.getTrackbarPos('frame',title)
                end_time = timing(current_frame)
                with open(annofile,'a') as f:
                    f.writelines(''.join([end_time,'\n']));
            elif wait_key == 46: # .
                clip_flag = 1
                tag_state = load_tag_state()
                try:
                    last_line = get_last_anno()
                except:
                    last_line = "NULL"
                try:
                    last_start,last_end = last_line.split('|')
                except:
                    last_start = last_line.split('|')[0]
                if last_start != 'NULL':
                    last_start_frame = get_position(last_start)
                else:
                    last_start_frame = position
                position = last_start_frame
                loop_flag = position
                media.set(cv2.CAP_PROP_POS_FRAMES,position)
            elif wait_key == 47: # /
                clip_flag = 1
                tag_state = load_tag_state()
                last_line = get_last_anno()
                try:
                    last_start,last_end = last_line.split('|')
                except:
                    last_end = last_line.split('|')[0]
                if last_end != 'NULL':
                    last_end_frame = get_position(last_end)
                else:
                    last_end_frame = position
                position = last_end_frame
                loop_flag = position
                media.set(cv2.CAP_PROP_POS_FRAMES,position)
            elif wait_key == 45: # -
                position = 1
                loop_flag = position
                media.set(cv2.CAP_PROP_POS_FRAMES,position)
            elif wait_key == 61: # =
                #if not load_tag_state() == 'None':
                #    position = frames
                #    loop_flag = position
                #    media.set(cv2.CAP_PROP_POS_FRAMES,position)
                position = frames
                loop_flag = position
                media.set(cv2.CAP_PROP_POS_FRAMES,position)

            success, frame = media.read()
            if success:
                frame_index += 1
                cv2.putText(frame,time_info,time_position,time_font,font_size,time_color,time_bold)
                cv2.putText(frame,help_info_1,help_position_1,help_font,font_size,help_color_1,help_bold_1)
                cv2.putText(frame,help_info_2,help_position_2,help_font,font_size,help_color_2,help_bold_2)
                cv2.putText(frame,help_info_3,help_position_3,help_font,font_size,help_color_2,help_bold_3)
                cv2.putText(frame,''.join(['[',str(serial_num),'/',str(total_num),']']),serial_position,help_font,font_size,time_color,help_bold_2)
                cv2.putText(frame,''.join([time_video_start]),time_video_start_position,font_size_time_video_start,font_bold_time_video_start,help_color_1,help_bold_1)
                if 'tag_state' in locals().keys():
                    cv2.putText(frame,tag_state,tag_state_position,tag_font,font_size,tag_state_color,tag_state_bold)
                if clip_flag == 1:
                    try:
                        last_tag = get_last_anno()
                    except:
                        last_tag = 'NULL'
                    cv2.putText(frame,last_tag,tag_position,tag_font,font_size,tag_color,tag_bold)
                cv2.imshow(title, frame)
                #if position == frames - 25:
                #    cv2.waitKey(0)

            else:
                cv2.destroyWindow(title)
                if edit_flag == 1:
                    cmd = ''.join(["vim ",annofile])
                    #os.system(cmd)
                if 'tag_state' in locals().keys():
                    if tag_state == 'xx':
                        mediafile_useless = Path(mediafile)
                        mediafile_useless.replace(''.join([mediafile,'.bak']))
                    if tag_state == 'ok':
                        mediafile_useless = Path(mediafile)
                        mediafile_useless.replace(''.join([mediafile,'.ok']))
                break
        print("frame_index = ", frame_index)
#lst_all_mediafile = ''.join(['ls -1th ',src(rootPath),'/*.mp4 > clip_dir.sh'])
#add_python_command = ''.join(['sed -i "s/^/python clip_single.py /p" ','clip_dir.sh'])
#clip_all_medias = ''.join(['bash clip_dir.sh'])
#os.system(lst_all_mediafile)
#os.system(add_python_command)
#os.system(clip_all_medias)

